#!/usr/bin/python
'''
From an annotated edge file it pulls the correctly mapped edges.
'''
import sys
import os

# parameters.
node_file = sys.argv[1]
edge_file = sys.argv[2]
test_file = sys.argv[3]

# Read in the node file.
print "Reading edge file."
fin = open(node_file, "rb")
lines = fin.readlines()
fin.close()

# Parse node file.
nodes = {}
print "Parsing node file."
for line in lines:
	# tokenize.
	tmp = line.strip().split("\t")
	
	# save id to chrom name.
	nodes[tmp[0]] = tmp

# Read in edge file.
print "Reading edge file."
fin = open(edge_file, "rb")
lines = fin.readlines()
fin.close()

# Parse edge file.
bad_pairs = set()
good_pairs = set()
print "Parsing edge file."
for line in lines:
	# tokenize.
	tmp = line.strip().split("\t")
	id1 = tmp[0]
	id2 = tmp[1]
	lstart1 = int(tmp[4])
	lstart2 = int(tmp[5])

	# sanity check.
	if id1 not in nodes:
		print "node error", id1
		sys.exit()	
	if id2 not in nodes:
		print "node error", id2
		sys.exit()

	# get global starts.
	gstart1 = int(nodes[id1][4]) + lstart1
	gstart2 = int(nodes[id2][4]) + lstart2
		
	# setup result to add.
	res = line
		
	# Check for inter chromosome pairs.
	if nodes[id1][3] != nodes[id2][3]:
		bad_pairs.add(res)
		continue
		
	# Check for greater than 5 stddev.
	gstart1 = int(nodes[id1][4]) + lstart1
	gstart2 = int(nodes[id2][4]) + lstart2
	
	if abs(gstart1 - gstart2) > ((5 * 250) + 1395):
		bad_pairs.add(res)
		continue
	
	if abs(gstart1 - gstart2) < ((5 * 250) - 1395):
		bad_pairs.add(res)
		continue
		
	# Must be a good pair.
	good_pairs.add(res)
	
# Read in predicted edge file.
print "Reading predicted edge file."
fin = open(test_file, "rb")
lines = fin.readlines()
fin.close()

	
# load predicted good edges.
pred_good = set()
print "Parsing predicted edge file."
for line in lines:
	# tokenize.
	tmp = line.strip().split("\t")
	id1 = tmp[0]
	id2 = tmp[1]
	lstart1 = int(tmp[4])
	lstart2 = int(tmp[5])

	# sanity check.
	if id1 not in nodes:
		print "node error", id1
		sys.exit()	
	if id2 not in nodes:
		print "node error", id2
		sys.exit()
		
	# setup result to add.
	res = line
		
	pred_good.add(res)


# calculate predicated bad edges.
print "calculating statistics."

# get predicted bad.
tot_edges = bad_pairs.union(good_pairs)
pred_bad = tot_edges.difference(pred_good)

# count number i said were good, but werent.
t1 = len( pred_good.difference(good_pairs) )

# count number i said were bad but were good.
t2 = len( pred_bad.difference(bad_pairs) )

# print results.
print "original edges\t%i" % (len(good_pairs) + len(bad_pairs))
print "filtered edges\t%i" % len(pred_good)
print "type1\t%i" % t1
print "type2\t%i" % t2

